from osgeo import gdal
import ogr
"""
    通过asc创建shp高程文件
    使用OGC创建shp文件数据源，获取OGR图层，打开DEM，最后在OGR图层上生成等高线
"""
# 高程DEM
source = "Grid/dem/dem.asc"
# 输出shp文件
output = "Grid/dem/contour"

ogr_driver = ogr.GetDriverByName("ESRI Shapefile")
ogr_ds = ogr_driver.CreateDataSource(output + ".shp")
# wkbLineString25D是包含几何高程值z的类型代码
ogr_lyr = ogr_ds.CreateLayer(output, geom_type=ogr.wkbLineString25D)
field_defn = ogr.FieldDefn("ID", ogr.OFTInteger)
ogr_lyr.CreateField(field_defn)
field_defn = ogr.FieldDefn("ELEV", ogr.OFTReal)
ogr_lyr.CreateField(field_defn)

ds = gdal.Open(source)
gdal.ContourGenerate(ds.GetRasterBand(1), 400, 10, [], 0, 0, ogr_lyr, 0, 1)
ogr_ds = None

"""
    通过PNGCanvas库绘制上述代码生成的等高线shp文件内容
"""
import shapefile
import pngcanvas

r = shapefile.Reader(output)  # output = "Grid/dem/contour" 该文件为上面代码生成的等高线文件

# 世界坐标到屏幕坐标转换，等价于 def world2screen(bbox, w, h, x, y):
xdist = r.bbox[2] - r.bbox[0]
ydist = r.bbox[3] - r.bbox[1]
width = 800
height = 600
xratio = width/xdist
yratio = height/ydist

# 遍历所有形状
contours = []
for shape in r.shapes():
    pixels = []
    for i in range(len(shape.parts)):
        pixels = []
        # 判断shape.points的点时候在范围内
        pt = None
        if i < len(shape.parts) - 1:
            pt = shape.points[shape.parts[i]:shape.parts[i+1]]
        else:
            pt = shape.points[shape.parts[i]:]
        # x, y 转换
        for x, y in pt:
            px = int(width - (r.bbox[2] - x) * xratio)
            py = int((r.bbox[3] - y) * yratio)
            pixels.append([px, py])
        contours.append(pixels)

# 创建输出绘制板
canvas = pngcanvas.PNGCanvas(width, height)
# PNGCanvas库可以接受颜色的rgba字节数组,此处定义为红色
red = [0xff, 0, 0, 0xff]
canvas.color = red
# 遍历多边形并绘制
for c in contours:
    canvas.polyline(c)
# 保存图片
with open(output+".png","wb") as f:
    f.write(canvas.dump())

# 用脚本添加dbf高程字段标记等高线的高程
# 学习网站https://www.usgs.gov/core-science-systems/eros/edna
r.close()
